from __future__ import absolute_import
import logging
#import time
import re
import os
import sys
try:
    if not os.environ.get('READTHEDOCS', None):
        from reparted import *
except Exception as e:
    print e
    print 'Make sure you installed the parted library.'
    sys.exit(-1)


from forest.api import Command, get_tool
logger = logging.getLogger('forest.commands.CreatePartitions')

from forest.resources.partitiontable import PartitionTable
class Partition(Command):
    ''' Creates all partitions on specified disks using partition-table resources.
    '''
    __register__ = 'commands.install.Partition'

    #requires_config = ['partitions', 'install_mode']

    def run(self, config):
        # Sometimes we dont want to partition. Skip if not defined.
        #if not config.get('partitions', False):
        #    return

        #installtree = config.get_install_root()
        self.mds, self.mddiskpartusage = self.get_disk_part_usage()
        print self.mds,  self.mddiskpartusage

        #diskspecs = config.get_type(Partitions, 'partitions')
        partitiontables = list(config.get_resources().iter_by_type('partition-table'))

        #diskspecs = PartitonTable
        if not  partitiontables:
            return


        # TODO: before creating partitions, make sure we are not installing on forests image (if it happens to be a disk or stick)
        # TODO: we need a way to rename the device in udev

        # stop and remove anything in raid so we dont get disk/partition in use errors
        for diskname, devs_parts in self.mddiskpartusage.items():
            for dev in devs_parts['devices']:
                mddev = '/dev/%s' % dev
                for line in get_tool('system.poll')(['mdadm', '--stop', mddev]):
                    print line


            for part in devs_parts['partitions']:
                devpart = '/dev/%s%s' % (diskname, part)
                for line in get_tool('system.poll')(['mdadm', '--misc', '--zero-superblock', devpart]):
                    print line

        # build the partitions per disk
        # commit after applying the partitions on a disk
        for partitiontable in partitiontables:
            for diskspec in partitiontable.get_disk_specs():

                diskid = diskspec.get_disk()
                diskname = '/dev/%s' % diskid

                if not config.is_dry_run():
                    dev = device.Device(diskname)

                Command.SIGNAL_COMMAND_PROGRESS(self, message='Partitioning device %s' % diskname)

                if not config.is_dry_run():
                    mydisk = disk.Disk(dev)
                    mydisk.delete_all()
                    mydisk.set_label('gpt')
                    Command.SIGNAL_COMMAND_PROGRESS(self, message='Some properties: name - {0}, length - {1}, block sizes - {2}'.format(diskname, dev.length, mydisk.block_sizes))
                last_start, last_end, last_length = 0, 0, 0

                total_created = 0
                for partition in diskspec.get_partitions():
                    part_size = partition.get_size()

                    if part_size == 'grow':
                        if not config.is_dry_run():
                            # minus double 34.. gpt copies
                            grow = dev.length - total_created-68#last minus to fudge a bug..
                            #print 'Growing', grow, total_created, total_created+grow, dev.length
                            sz = size.Size(length=(grow*(1000**2))/ dev.sector_size, units='MB', dev=dev)
                            sz.sectors = grow
                    else:
                        if not config.is_dry_run():
                            sz = size.Size(length=part_size, units='MB', dev=dev)

                    if not config.is_dry_run():
                        new = disk.Partition(mydisk, sz, fs=partition.get_fs(), align='minimal')
                        mydisk.add_partition(new)
                        Command.SIGNAL_COMMAND_PROGRESS(self, message='Added partition: geom - {0} , disk - {1}'.format(str(new.geom), str(new.disk)) )

                        for flag, value in partition.get_flags().items():
                            new.set_flag(flag.upper(), value)

                        last_start, last_end, last_length = new.geom
                        total_created += last_length
                if not config.is_dry_run():
                    mydisk.commit()

        #time.sleep(20)
        # Dunno, sometimes this is needed for things to get reconized (udev settle?)
        get_tool('parted.probe')
        get_tool('udev.settle')
        return


    @staticmethod
    def get_disk_part_usage():
        ''' Returns a dictionary of md\d+ -> {'devices':[],'partitions':[]}

            where the sub dict devices is a list of disk devices belonging to the raid,
            and partitions is a list diskdevices with partition numbers
        '''
        mddiskpartusage = {}
        mds = set()
        raidmatch = re.compile('md\d+')
        diskmatch = re.compile('([vhms]d[a-z])(\d+)')

        for dev in os.listdir('/sys/block'):
            match = raidmatch.match(dev)

            if match is None:
                continue

            mds.add(match.group(0))

            for slave in os.listdir('/sys/block/%s/slaves' % match.group(0)):

                diskm = diskmatch.match(slave)

                if diskm is None:
                    continue

                #diskpartkey = diskm.group(0)
                diskname  = diskm.group(1)
                if diskname not in mddiskpartusage:
                    mddiskpartusage[diskname] = {'devices':[], 'partitions':[]}

                mddiskpartusage[diskname]['devices'].append(dev)
                mddiskpartusage[diskname]['partitions'].append(diskm.group(2))


        return mds, mddiskpartusage

if __name__ == '__main__':

    c = Partition({})
    print c.get_disk_part_usage()
